# -*- coding: utf-8 -*-
#!/usr/bin/env python
# view.py
# 用于可视化四节点壳单元有限元结果
# 用法： python view.py 结果文件路径 变量名 --scale 变形放大倍数
import numpy as np
import pyvista as pv
import argparse
from pathlib import Path

def read_fem_file(filepath)->tuple[np.ndarray,np.ndarray,np.ndarray,np.ndarray]:
    """
    读取有限元输出文件，返回节点、单元、位移、反力
    """
    with open(filepath, 'r') as f:
        lines = f.readlines()

    nodes = []
    elements = []
    displacements = []
    reactions = []

    current_section = None

    for line in lines:
        line = line.strip()
        if line.startswith('#'):
            if 'Node Coordinates' in line:
                current_section = 'nodes'
            elif 'Element Connectivity' in line:
                current_section = 'elems'
            elif 'Displacements' in line:
                current_section = 'disp'
            elif 'Reactions' in line:
                current_section = 'reactions'
            elif 'End' in line:
                current_section = None
            continue
        
        if current_section == 'nodes':
            parts = line.split()
            x, y, z = map(float, parts[1:4])
            nodes.append([x, y, z])

        elif current_section == 'elems':
            parts = line.split()
            conn = list(map(int, parts[1:]))  # 四节点
            elements.append(conn)

        elif current_section == 'disp':
            parts = line.split()
            ux, uy, uz = map(float, parts[1:4])
            displacements.append([ux, uy, uz])

        elif current_section == 'reactions':
            parts = line.split()
            if not parts:
                continue
            rx, ry, rz = map(float, parts[1:4])
            reactions.append([rx, ry, rz])
    nodes=np.array(nodes,dtype=float)
    elements=np.array(elements,dtype=int)-1  # 转换为0索引
    displacements=np.array(displacements,dtype=float)
    reactions=np.array(reactions,dtype=float)
    return nodes, elements, displacements, reactions


def create_mesh(nodes, elements:np.ndarray)->pv.UnstructuredGrid:
    """
    使用 pyvista 创建网格（四边形单元）
    """
    points = nodes.copy()
    cells = []
    for i in range(elements.shape[0]):
        conn= elements[i,:]
        # PyVista: [n, v0, v1, v2, v3]，n=4 表示四边形
        cell = [8] + conn.tolist()
        cells.append(cell)

    cells = np.array(cells,dtype=int)
    mesh = pv.UnstructuredGrid(cells, [12] * np.ones(shape=(cells.shape[0],),dtype=int), points)
    return mesh


def main():
    parser = argparse.ArgumentParser(description="可视化四节点壳单元有限元结果")
    parser.add_argument("file", type=str, help="有限元结果文件路径")
    parser.add_argument("variable", type=str, choices=['Ux', 'Uy', 'Uz', 'Rx', 'Ry', 'Rz'],
                        help="要可视化的变量 (Ux, Uy, Uz, Rx, Ry, Rz)")
    parser.add_argument("--scale", type=float, default=1.0, help="变形放大倍数 (默认: 1.0)")

    args = parser.parse_args()

    filepath = Path(args.file)
    if not filepath.exists():
        print(f"错误: 文件 {filepath} 不存在！")
        return

    # 读取数据
    print(f"正在读取文件: {filepath}")
    nodes, elements, displacements, reactions = read_fem_file(filepath)

    # 创建网格
    mesh = create_mesh(nodes, elements)

    # 提取变量数据
    var_map = {
        'Ux': (displacements, 0, "Displacement X"),
        'Uy': (displacements, 1, "Displacement Y"),
        'Uz': (displacements, 2, "Displacement Z"),
        'Rx': (reactions, 0, "Reaction X"),
        'Ry': (reactions, 1, "Reaction Y"),
        'Rz': (reactions, 2, "Reaction Z"),
    }

    data_dict, idx, label = var_map[args.variable]
    scalar_data = data_dict[:, idx]  # 提取对应的变量数据

    # 添加标量到网格
    mesh[args.variable] = scalar_data
    # 不显示变形
    plotter = pv.Plotter()
    plotter.add_mesh(mesh, scalars=scalar_data, show_edges=True, cmap='viridis')
    plotter.add_title(label)

    # plotter.show_grid()
    plotter.show()


if __name__ == "__main__":
    main()